import xml.etree.ElementTree as ET
from datetime import datetime

from base import claude3_api as chat_ai
from model.models import db, TSysNovel, TSysCharacterAppearance, TSysNovelChapter, TSysNovelOutline

pic_prompt_style = "动画风格，环境为[historical_background]的风格。"
character_prompt_style = "人物外表为[historical_background]的人外表和服饰, 不要出现其他年代人物外表和服饰。"

def map_historical_background(code):
    if code == 1:
        return "现代"
    elif code == 2:
        return "当代"
    elif code == 3:
        return "近代"
    elif code == 4:
        return "古代"
    elif code == 5:
        return "古代架空"
    elif code == 6:
        return "未来"
    else:
        return "现代"

def character_appearance_extraction(novel_id):
    if novel_id is None:
        return {"success": False, "value": None, "msg": "novel_id为空"}
    role = ""

    novel = TSysNovel.query.filter_by(id=novel_id).first()
    if novel is None:
        return {"success": False, "value": None, "msg": "novel为空"}

    user_id = novel.user_id
    historical_background ,role = get_historical_background_and_role(novel=novel)
    replaced_character_prompt_style = str(character_prompt_style).replace("[historical_background]", map_historical_background(code=historical_background))

    #读取视频剧本大纲
    character_appearance_extraction_promot = "以下是视频剧本的各个角色，请分别给出他们的样貌描述，描述需要精确简短并且可以作为用来生成图片的提示词，描述包含：名字、 性别、 年龄、头发、 五官、上衣、下衣。每个人物的样貌以<role>开始，以</role>结束。" + replaced_character_prompt_style #人物外貌提取提示词
    message = []
    message.append({"role": "user","content": str(character_appearance_extraction_promot + role)})
    character_appearance = ""
    answer = None
    try:
        answer = chat_ai.chat(message)
    except Exception as e:
        print("AI 回答异常, 信息为:" + e)
    if answer is None:
        return {"success": False, "value": None, "msg": "answer为空"}

    #抽取人物样貌
    roles_xml = answer.get("answer")
    tree = ET.fromstring("<roles>" + roles_xml + "</roles>")
    roles = tree.findall("role")
    if roles is None:
        return {"success": False, "value": None, "msg": "roles解析错误"}

    for role in roles:
        # character_appearance += replaced_character_prompt_style + role.text
        character_appearance += role.text

    is_exist_character_appearance = get_character_appearance(novel_id=novel_id)
    if is_exist_character_appearance is None or is_exist_character_appearance.get("success") == False:
        characterAppearance = TSysCharacterAppearance(
            user_id=user_id,
            correlation_id=novel_id,
            correlation_type='novel',
            character_appearance=character_appearance,
            create_date=datetime.now(),
            updated_date=datetime.now()
        )
        db.session.add(characterAppearance)
        db.session.commit()
    elif is_exist_character_appearance.get("success") == True:
        character_appearances = TSysCharacterAppearance.query.filter_by(correlation_id=novel_id, correlation_type='novel').first()
        character_appearances.character_appearance = character_appearance
        db.session.commit()
    else:
        return {"success": False, "value": None, "msg": "创建人物外貌失败,未保存成功或未存在需要修改的样貌数据"}

    return {"success": True, "value": None, "msg": "创建人物外貌成功"}

def create_playbook_by_chapter(chapter_id, character_appearances):
    if chapter_id is None:
        print("chapter_id为空")
        return {"success": False, "value": None, "msg": "chapter_id为空"}

    novel_chapter = TSysNovelChapter.query.filter_by(id=chapter_id).first()
    if novel_chapter is None:
        print("章节不存在")
        return {"success": False, "value": None, "msg": "章节不存在"}

    charter_content = novel_chapter.content
    if charter_content is None:
        print("章节内容为空")
        return {"success": False, "value": None, "msg": "章节内容为空"}

    vedio_playbook_promopt_1 = "标签‘<content>’和标签‘</content>’中是一个篇小说的内容，我们希望将其用图片呈现。 返回值为XML形式。有唯一的一个根节点：根节点以<playbooks>开始，以</playbooks>结束。"
    vedio_playbook_promopt_2 = "1，将小说原文整篇内容分割为多个分镜段落，段落数量{必须涵盖}小说原文所有内容，需要包含在返回中，每个段落以<para>开始，以</para>结束。中间包含段落内容，如：<para>这里是段落内容</para>。"
    vedio_playbook_promopt_3 = "2，给以上的每个分镜段落生成图片的提示词，一个段落包含不能少于1张图片。每一个图片的提示词以<pic_promopt>开始，以</pic_promopt>结束， 中间包含段落内容，如：<pic_promot>图画的提示词</pic_promot>。涉及的人物外貌描述为：" + character_appearances + "，生成图片的提示词中{所有}人物{必须}使用人物外貌描述中的对应内容，人物{必须为动画风格}。提示词必须简洁明了，使用中文输出。一个分镜段落的所有<pic_promot>，包含在<pic_promopts>开头，以</pic_promopt>结尾的结构中。"
    # vedio_playbook_promopt_4 = "3，给每个图片生成字幕，字幕数量和图片数量相同，每段字幕为对应的段落内容，不能修改内容，每个字幕以<subtitle>开始，以</subtitle>结束，全部字幕内容包含在<subtitles></subtitles>中。返回格式如：<subtitles><subtitle>字幕内容1</subtitle><subtitle>字幕内容2</subtitle></subtitles>。"
    vedio_playbook_promopt_5 = "3，我们来约定返回格式。每个段落是一个分镜，以<playbook>开始，以</playbook>结束，每个playbook包含一个<para>，它对应的图画提示词<pic_promopts>。"  # 以及字幕<subtitles>

    promopt = vedio_playbook_promopt_1 + vedio_playbook_promopt_2 + vedio_playbook_promopt_3 + vedio_playbook_promopt_5 + "小说原文内容如下： <content>" + charter_content + "</content>"
    #章节分镜提示词
    message = []
    message.append({"role": "user","content": promopt})

    playbooks_json = []
    try:
        answer = chat_ai.chat(message)
        if answer is None:
            print("AI回答内容为空")
            return {"success": False, "value": None, "msg": "AI回答内容为空"}

        # 抽取人物样貌
        playbooks_xml = answer.get("answer")

        if playbooks_xml is None:
            return {"success": False, "value": None, "msg": "解析AI回答错误"}

        if playbooks_xml.startswith("<playbooks") == False:
            playbooks_start_index = playbooks_xml.index("<playbooks>")
            playbooks_end_index = playbooks_xml.index("</playbooks>")
            playbooks_xml = playbooks_xml[playbooks_start_index:playbooks_end_index + 12]

        tree = ET.fromstring(playbooks_xml)
        if tree is None:
            return {"success": False, "value": None, "msg": "解析AI回答错误"}

        playbooks = tree.findall("playbook")
        if playbooks is None:
            return {"success": False, "value": None, "msg": "解析AI回答错误"}

        novel = TSysNovel.query.filter_by(id=novel_chapter.novel_id).first()
        historical_background, role = get_historical_background_and_role(novel=novel)
        replaced_pic_prompt_style = str(pic_prompt_style).replace("[historical_background]",map_historical_background(code=historical_background))

        pic_prompts_str = ''
        playboosks_json = []
        for playbook in playbooks:
            playbook_json={}
            subtitles = []
            if playbook.find("para") is not None:
                playbook_json["para"] = playbook.find("para").text
                subtitles.append(playbook.find("para").text)

            playbook_json["subtitles"] = subtitles

            pic_prompts_json = []
            pic_prompts = playbook.find("pic_promopts")
            if playbook.find("pic_promopts") is not None:
                for pic_prompt in pic_prompts:
                    pic_prompt_json = {}
                    pic_prompt_json["pic_promopt"] = "生成的图片风格{必须为}" + replaced_pic_prompt_style + "图片内容为:" + pic_prompt.text
                    pic_prompts_str +=  pic_prompt_json["pic_promopt"] + '。'
                    pic_prompts_json.append(pic_prompt_json)

            playbook_json["pic_promopts"] = pic_prompts_json
            playboosks_json.append(playbook_json)

        print(pic_prompts_str)
    except Exception as e:
        print(e)
        return {"success": False, "value": None, "msg": "生成剧本失败"}

    return {"success": True, "value": playboosks_json, "msg": "生成剧本成功"}


def get_character_appearance(novel_id):
    if novel_id is None:
        return {"success": False, "value": None, "msg": "小说id为空"}

    character_appearances = TSysCharacterAppearance.query.filter_by(correlation_id=novel_id,correlation_type='novel').first()
    if character_appearances is not None:
        return {"success": True, "value": character_appearances.character_appearance, "msg": "use_id为空"}
    else:
        return {"success": False, "value": None, "msg": "人物样貌为空"}

def get_historical_background_and_role(novel):
    if novel is None:
        return "", ""

    outline_id = novel.outline_id
    outline = TSysNovelOutline.query.filter_by(id=outline_id).first()
    if outline is None:
        return False

    role = outline.character_settings
    historical_background = outline.historical_background
    return historical_background, role


#版本2方法
def create_playbook_by_chapter_v2(chapter_id, character_appearances):
    if chapter_id is None:
        print("chapter_id为空")
        return {"success": False, "value": None, "msg": "chapter_id为空"}

    novel_chapter = TSysNovelChapter.query.filter_by(id=chapter_id).first()
    if novel_chapter is None:
        print("章节不存在")
        return {"success": False, "value": None, "msg": "章节不存在"}

    charter_content = novel_chapter.content
    if charter_content is None:
        print("章节内容为空")
        return {"success": False, "value": None, "msg": "章节内容为空"}
    #{必须涵盖}小说原文所有内容，
    vedio_playbook_promopt_1 = "标签‘<content>’和标签‘</content>’中是一个篇小说的内容，我们希望将其用视频的形式呈现。 返回值为XML形式。有唯一的一个根节点：根节点以<playbooks>开始，以</playbooks>结束。"
    vedio_playbook_promopt_2 = "1，将小说原文整篇内容分割为多个分镜，分镜数量{不多于}6段，需要包含在返回中，每个分镜以<para>开始，以</para>结束。中间包含分镜内容，如：<para>这里是分镜内容</para>。"
    vedio_playbook_promopt_3 = "2，给以上的每个分镜段落生成视频的提示词。每个段落生成{6}个视频提示词。每一个视频的提示词以<video_promopt>开始，以</video_promopt>结束， 中间视频提示词内容，如：<video_promopt>视频的提示词</video_promopt>。涉及的人物外貌描述为：" + character_appearances + "，生成图片的提示词中{所有}人物{必须}使用人物外貌描述中的对应内容，人物{必须为动画风格}。提示词必须简洁明了，使用中文输出。一个分镜段落的所有<video_promopt>，包含在<video_promopts>开头，以</video_promopts>结尾的结构中。"
    # vedio_playbook_promopt_4 = "3，给每个图片生成字幕，字幕数量和图片数量相同，每段字幕为对应的段落内容，不能修改内容，每个字幕以<subtitle>开始，以</subtitle>结束，全部字幕内容包含在<subtitles></subtitles>中。返回格式如：<subtitles><subtitle>字幕内容1</subtitle><subtitle>字幕内容2</subtitle></subtitles>。"
    vedio_playbook_promopt_5 = "3，我们来约定返回格式。每个段落是一个分镜，以<playbook>开始，以</playbook>结束，每个playbook包含一个<para>，它对应的视频提示词<video_promopts>。"  # 以及字幕<subtitles>

    promopt = vedio_playbook_promopt_1 + vedio_playbook_promopt_2 + vedio_playbook_promopt_3 + vedio_playbook_promopt_5 + "小说原文内容如下： <content>" + charter_content + "</content>"
    #章节分镜提示词
    message = []
    message.append({"role": "user","content": promopt})

    playbooks_json = []
    try:
        answer = chat_ai.chat(message)
        if answer is None:
            print("AI回答内容为空")
            return {"success": False, "value": None, "msg": "AI回答内容为空"}

        # 抽取人物样貌
        playbooks_xml = answer.get("answer")

        if playbooks_xml is None:
            return {"success": False, "value": None, "msg": "解析AI回答错误"}

        if playbooks_xml.startswith("<playbooks") == False:
            playbooks_start_index = playbooks_xml.index("<playbooks>")
            playbooks_end_index = playbooks_xml.index("</playbooks>")
            playbooks_xml = playbooks_xml[playbooks_start_index:playbooks_end_index + 12]

        tree = ET.fromstring(playbooks_xml)
        if tree is None:
            return {"success": False, "value": None, "msg": "解析AI回答错误"}

        playbooks = tree.findall("playbook")
        if playbooks is None:
            return {"success": False, "value": None, "msg": "解析AI回答错误"}

        novel = TSysNovel.query.filter_by(id=novel_chapter.novel_id).first()
        historical_background, role = get_historical_background_and_role(novel=novel)
        replaced_pic_prompt_style = str(pic_prompt_style).replace("[historical_background]",map_historical_background(code=historical_background))

        video_prompts_str = ''
        playboosks_json = []
        for playbook in playbooks:
            playbook_json={}
            subtitles = []
            if playbook.find("para") is not None:
                playbook_json["para"] = playbook.find("para").text
                subtitles.append(playbook.find("para").text)

            playbook_json["subtitles"] = subtitles

            video_prompts_json = []
            video_prompts = playbook.find("video_promopts")
            if playbook.find("video_promopts") is not None:
                for video_prompt in video_prompts:
                    video_prompt_json = {}
                    video_prompt_json["video_promopt"] = "生成的视频风格{必须为}" + replaced_pic_prompt_style + "视频内容为:" + video_prompt.text
                    video_prompts_str +=  video_prompt_json["video_promopt"] + '。'
                    video_prompts_json.append(video_prompt_json)

            playbook_json["video_promopts"] = video_prompts_json
            playboosks_json.append(playbook_json)

    except Exception as e:
        print(e)
        return {"success": False, "value": None, "msg": "生成剧本失败"}

    return {"success": True, "value": playboosks_json, "msg": "生成剧本成功"}